import heapq
from graphviz import Graph
import networkx as nx

import matplotlib.pyplot as plt


def create_graph():
    # 定义图
    graph = {
        'A': {'B': 310, 'C': 695, 'D': 120},
        'B': {'A': 310, 'C': 450},
        'C': {'A': 695, 'B': 450, 'G': 1130, 'F': 1420},
        'D': {'A': 120, 'E': 600},
        'E': {'C': 200, 'D': 700, "G": 500},
        'G': {'C': 1130, 'F': 115},
        'F': {'C': 1420, 'G': 115}
    }

    return graph

'''
第四题：实现Dijkstra算法和A*算法
'''
def dijkstra(graph, start, end):
    # 节点距离字典，
    distances = {node: float('inf') for node in graph}
    # 将起始点的距离设为0
    distances[start] = 0
    # 前置节点字典
    parents = {node: None for node in graph}
    # 创建空堆，并将起始节点压入堆中
    pq = [(0, start)]
    # 遍历堆 求解图中从源节点到其他节点的最短路径
    while len(pq) > 0:
        # 从堆中弹出当前最小的节点 返回距路及当前节点
        (dis, node) = heapq.heappop(pq)
        # 如果当前节点已经访问过，直接跳过
        if dis > distances[node]:
            continue
        # 重新构成一个邻接表
        graphs = graph[node].items()
        # 遍历当前节点的所有邻居
        for neighbor, weight in graphs:
            # 计算出从源节点到邻居节点的距离
            new_distance = dis + weight
            # 如果新的距离比已知的距离小，则更新距离和前置节点
            if new_distance < distances[neighbor]:
                # 如果新的距离new_distance小于之前计算的最短距离distances[neighbor]，则更新距离和父节点信息，并将邻居节点加入最小堆pq中
                # 更新了从源节点到邻居节点的最短距离
                distances[neighbor] = new_distance
                # 更新了邻居节点的父节点为当前节点
                parents[neighbor] = node
                # 将邻居节点压入堆中
                heapq.heappush(pq, (new_distance, neighbor))

    # 根据前置节点字典构建最短路径
    path = []
    node = end
    # 构造从起点到终点的最短路径
    while node != start:
        # 从终点开始，将每个节点依次插入到路径列表 path的第一个位置
        path.insert(0, node)
        # 将当前节点更新为其父节点，继续向前回溯，直到回溯到起点
        node = parents[node]
    # 将起点插入到路径列表的第一个位置
    path.insert(0, start)
    # 返回最短路径和距离
    return (path, distances[end])


def astar(graph, start, goal):
    # 初始化起点和终点
    # 定义优先队列
    pQueue = [(0, start)]
    came_from = {}
    cost_so_far = {start: 0}

    # A*算法主循环
    while pQueue:
        # 获取当前最小的节点
        current_cost, current_node = heapq.heappop(pQueue)

        # 如果当前节点是目标节点，则算法结束
        if current_node == goal:
            break

        # 遍历当前节点的所有邻居
        for next_node, cost in graph[current_node].items():
            # 计算从起点到下一个节点的预估距离
            priority = cost + cost_so_far[current_node]

            # 如果下一个节点还没有被探索过，或者到达它的路径更短
            if next_node not in cost_so_far or priority < cost_so_far[next_node]:
                # 更新到达下一个节点的代价和路径
                cost_so_far[next_node] = priority
                came_from[next_node] = current_node
                heapq.heappush(pQueue, (priority, next_node))

    # 根据 came_from 字典生成最短路径
    path = [goal]
    while path[-1] != start:
        path.append(came_from[path[-1]])
    #反转列表中的元素顺序
    path.reverse()
    # 返回最短路径和路径长度
    return (path, cost_so_far[goal])


def tsxtA():
    graph = create_graph()
    path,distance = astar(graph, 'A', 'F')
    print('A * 最短路径：', '->'.join(path), distance)


def testdijkstra():
    graph = create_graph()
    path, distance = dijkstra(graph, 'A', 'F')
    print('dijkstra最短路径：', '->'.join(path), distance)

    # 测试代码

    # 创建一个有向图对象 G
    G = nx.DiGraph()
    # 遍历获取前节点
    for beforkey in graph:
        # 将元组中的邻居节点和权重值分别赋值给变量 neighbor 和 weight
        for neighbor, weight in graph[beforkey].items():
            # 向一个无向图中添加一条带权重的边
            G.add_edge(beforkey, neighbor, weight=weight)
    # 计算节点布局
    pos = nx.spring_layout(G)

    # 从带权有向图 G 中提取每条边的权重，并将其保存到字典 edge_labels 中

    _dict = {}
    for u, v, d in G.edges(data=True):
        _dict[(u, v)] = d['weight']

    # 绘制图中边的权重
    nx.draw_networkx_edge_labels(G, pos, edge_labels=_dict, font_size=10)

    # G 要可视化的图
    # pos 节点 ID 映射到它们的 (x, y) 坐标的字典
    # with_labels 是否显示节点标签
    # node_color 节点的颜色
    # edge_color 边的颜色
    nx.draw(G, pos, with_labels=True, node_color='lightblue', edge_color='grey',
            width=1, alpha=0.7, font_size=12)
    # 遍历路径
    for i in range(len(path) - 1):
        # 绘制节点的边和权
        nx.draw_networkx_edges(G, pos, edgelist=[(path[i], path[i + 1])], edge_color='blue', width=2)
    # 生成领结矩阵
    plt.show()


if __name__ == '__main__':
    tsxtA()
    testdijkstra()
